# -*- coding: utf-8 -*-
"""
1. BaseCollection 类
主要作用：
作为基础容器类，管理一组相关实例
提供基本的集合操作功能：添加、获取、遍历和存储这些实例
实现了集合的基本功能（迭代、索引、长度计算）
提供计算所有包含元素的总边界框的能力
核心功能：
__getitem__, __iter__, __len__：实现类似列表的基本操作
append, extend：添加元素的方法
reset：重置集合内容
store：将集合内容转为JSON格式
bbox：计算包含所有实例的最小边界框
2. Collection 类
主要作用：
继承自BaseCollection，并添加了IText接口
专注于元素的分组和排序功能
实现了各种基于空间关系的分组算法
提供考虑阅读顺序的排序功能
核心功能：
group：根据用户定义的标准对实例进行分组
group_by_columns, group_by_rows, group_by_physical_rows：按列、行、物理行分组
group_by_connectivity：基于元素的连接性进行分组
sort_in_reading_order, sort_in_line_order：按阅读顺序或行顺序排序
sort_in_reading_order_plus：增强版阅读顺序排序
3. ElementCollection 类
主要作用：
继承自Collection，专门用于管理Element实例
在处理Element时维护父子关系
提供布局分析相关的方法
支持基于空间关系的实例筛选
核心功能：
重写的append、insert方法：处理Element时更新父级边界框
is_flow_layout：判断包含的元素是否为流式布局
contained_in_bbox：过滤包含在目标边界框中的实例
split_with_intersection：将实例分为相交和不相交两组
1. 页面元素组织结构的基础
这些Collection类为整个pdf2docx项目提供了基础的数据组织结构，用于管理各类PDF元素：
Pages类（继承自BaseCollection）：管理文档的所有页面
Shapes类（继承自ElementCollection）：管理图形元素如线条、填充区域、超链接
Lines类（继承自ElementCollection）：管理文本行集合
Cells类（继承自ElementCollection）：管理表格单元格
Rows类（继承自ElementCollection）：管理表格行
2. PDF元素的空间分析
在PDF解析过程中，这些类提供了关键的空间关系分析功能：
文本排版分析：利用group_by_columns、group_by_rows等方法识别文本的列和行
阅读顺序确定：使用sort_in_reading_order等方法确定文本的阅读顺序
表格结构识别：利用元素的空间关系分析表格结构（行、列、单元格）
元素分组：使用group和group_by_connectivity方法将相关元素组织在一起
3. 处理元素间的层次关系
ElementCollection类特别关注元素间的父子关系：
当添加元素时自动设置父级（如将Cell关联到Table）
更新父级的边界框，确保父级容器能完全包含所有子元素
提供基于空间关系的选择和过滤功能
"""
"""实例集合，例如 Blocks, Lines, Spans, Shapes 等。"""

import fitz
from .Element import Element
from .share import IText, TextDirection
from .algorithm import solve_rects_intersection, graph_bfs

# BaseCollection 类是一个基础容器类，用于管理一组相关的 Element 实例，它提供了基本的集合操作功能，例如添加、获取、遍历和存储这些实例
class BaseCollection:
    """基础集合类，表示实例列表。"""

    def __init__(self, instances: list = None, parent=None):
        """从实例列表初始化集合。"""
        self._parent = parent
        self._instances = []
        self.extend(instances or [])  # 注意：默认排除空实例

    def __getitem__(self, idx):
        # 获取集合中的第idx个实例
        # 如果idx超出范围，则抛出IndexError
        # 否则返回第idx个实例
        try:
            instances = self._instances[idx]
        except IndexError:
            msg = f"集合索引 {idx} 超出范围。"
            raise IndexError(msg)
        else:
            return instances
        
    def __iter__(self):
        # 支持迭代
        return (instance for instance in self._instances)

    def __len__(self):
        # 返回集合中的实例数量
        return len(self._instances)

    @property
    def parent(self):
        # 返回集合的父级
        return self._parent

    @property
    def bbox(self):
        """组合集合的边界框。"""
        rect = fitz.Rect()
        for instance in self._instances:
            # Rect中的|实现了运算符重载
            # 例如rect1 | rect2返回的是能够包裹这两个矩形的最小矩形rect3
            # 函数这里遍历了所有的instance，所以是获取一个能够包裹所有instance实例的最小的矩形
            rect |= instance.bbox
        return fitz.Rect(
            [round(x, 1) for x in rect]
        )  # 注意：四舍五入以避免数字误差

    def append(self, instance):
        # 附加一个实例
        if not instance:
            return
        self._instances.append(instance)

    def extend(self, instances: list):
        # 扩展集合，添加多个实例
        if not instances:
            return
        for instance in instances:
            self.append(instance)

    def reset(self, instances: list = None):
        """重置实例列表。
        用新的实例列表替换当前实例列表，
        如果实例列表为空，则清空当前实例列表
        
        参数：
            instances (list, optional): 重置为目标实例。默认为 None。

        返回：
            BaseCollection: 自身
        """
        self._instances = []
        self.extend(instances or [])
        return self

    def store(self):
        """以 json 格式存储属性。"""
        return [instance.store() for instance in self._instances]

    def restore(self, *args, **kwargs):
        """从字典列表构建集合。"""
        raise NotImplementedError


class Collection(BaseCollection, IText):
    """实例集合，专注于元素的分组和排序。"""

    @property
    def text_direction(self):
        """
        获取文本方向。所有实例必须具有相同的文本方向。
        
        获取所有实例的文本方向，去重，如果只有一个方向，则返回该方向，否则返回TextDirection.MIX
        """
        res = set(instance.text_direction for instance in self._instances)
        return list(res)[0] if len(res) == 1 else TextDirection.MIX

    def group(self, fun):
        """根据用户定义的标准对实例进行分组。

        参数：
            fun (function): 带有2个参数表示2个实例(Element)并返回布尔值的函数。

        返回：
            list: 分组后的 ``Collection`` 实例列表。

        示例 1::

            # 对相互交叉的实例进行分组
            fun = lambda a,b: a.bbox & b.bbox

        示例 2::

            # 对水平对齐的实例进行分组
            fun = lambda a,b: a.horizontally_aligned_with(b)

        .. 注意::
            这等同于一个图搜索问题，构建邻接列表，然后搜索图以查找所有连接的组件。
        """
        # 构建邻接列表：
        # 第i项是一组索引，它们与第i个实例相连。
        # 注意: O(n^2)方法，但当n<1000时是可接受的(~0.2s)，这满足于页块
        num = len(self._instances)
        index_groups = [set() for i in range(num)]  # type: list[set]
        for i, instance in enumerate(self._instances):
            # 当前实例与其后所有实例的连接
            for j in range(i + 1, num):
                if fun(instance, self._instances[j]):
                    index_groups[i].add(j)
                    index_groups[j].add(i)

        # 搜索图 -> 实例的分组索引
        groups = graph_bfs(index_groups)
        groups = [
            self.__class__([self._instances[i] for i in group]) for group in groups
        ]
        return groups

    def group_by_connectivity(self, dx: float, dy: float):
        """将连接的实例收集到同一组中。

        参数：
            dx (float): 定义连接性的x容差
            dy (float): 定义连接性的y容差

        返回：
            list: 分组后的 ``Collection`` 实例列表。

        .. 注意::
            * 这等同于一个图遍历问题，在构建邻接列表时尤为关键，尤其是大量顶点(路径)的情况。

            * 检查路径之间的交叉实际上是一个矩形交叉问题，已在许多文献中进行了研究。
        """
        # 构建图 -> 邻接列表：
        # 第i项是一组索引，它们与第i个实例相连
        num = len(self._instances)
        index_groups = [set() for _ in range(num)]  # type: list[set]

        # 解决矩形交叉问题
        i_rect_x, i = [], 0
        d_rect = (-dx, -dy, dx, dy)
        for rect in self._instances:
            points = [a + b for a, b in zip(rect.bbox, d_rect)]  # 考虑容差
            i_rect_x.append((i, points, points[0]))
            i_rect_x.append((i + 1, points, points[2]))
            i += 2
        i_rect_x.sort(key=lambda item: item[-1])
        solve_rects_intersection(i_rect_x, 2 * num, index_groups)

        # 搜索图 -> 实例的分组索引
        groups = graph_bfs(index_groups)
        groups = [
            self.__class__([self._instances[i] for i in group]) for group in groups
        ]
        return groups

    def group_by_columns(
        self, factor: float = 0.0, sorted: bool = True, text_direction: bool = False
    ):
        """根据边界框将元素分组为列。"""
        # 按列分割
        fun = lambda a, b: a.vertically_align_with(
            b, factor=factor, text_direction=text_direction
        )
        groups = self.group(fun)

        # 如果排序，则按x方向增加
        if sorted:
            idx = 3 if text_direction and self.is_vertical_text else 0
            groups.sort(key=lambda group: group.bbox[idx])

        return groups

    def group_by_rows(
        self, factor: float = 0.0, sorted: bool = True, text_direction: bool = False
    ):
        """根据边界框将元素分组为行。"""
        # 按行分割
        fun = lambda a, b: a.horizontally_align_with(
            b, factor=factor, text_direction=text_direction
        )
        groups = self.group(fun)

        # 如果排序，则按y方向增加
        if sorted:
            idx = 0 if text_direction and self.is_vertical_text else 1
            groups.sort(key=lambda group: group.bbox[idx])

        return groups

    def group_by_physical_rows(
        self, sorted: bool = False, text_direction: bool = False
    ):
        """将行分组为物理行。"""
        fun = lambda a, b: a.in_same_row(b)
        groups = self.group(fun)

        # 如果排序，则按y方向增加
        if sorted:
            idx = 0 if text_direction and self.is_vertical_text else 1
            groups.sort(key=lambda group: group.bbox[idx])

        return groups

    def sort_in_reading_order(self):
        """按阅读顺序（考虑文本方向）对集合实例进行排序，
        例如，对于正常阅读方向：从上到下，从左到右。
        """
        if self.is_horizontal_text:
            self._instances.sort(key=lambda e: (e.bbox.y0, e.bbox.x0, e.bbox.x1))
        else:
            self._instances.sort(key=lambda e: (e.bbox.x0, e.bbox.y1, e.bbox.y0))
        return self

    def sort_in_line_order(self):
        """在考虑文本方向的物理行中对集合实例进行排序，
        例如，对于正常阅读方向：从左到右。
        """
        if not self.is_vertical_text:
            self._instances.sort(key=lambda e: (e.bbox.x0, e.bbox.y0, e.bbox.x1))
        else:
            self._instances.sort(key=lambda e: (e.bbox.y1, e.bbox.x0, e.bbox.y0))
        return self

    def sort_in_reading_order_plus(self):
        """按阅读顺序对实例进行排序，特别是对于同一行中的实例。
        以自然阅读方向为例：行的阅读顺序，同一行中的实例从左到右。
        在下面的例子中，A 在 B 之前::

                         +-----------+
            +---------+  |           |
            |   A     |  |     B     |
            +---------+  +-----------+

        步骤：

            * 按阅读顺序对元素排序，即从上到下，从左到右。
            * 将元素分组到行中。
            * 对行中的元素排序：从左到右。
        """
        instances = []
        for row in self.group_by_physical_rows(sorted=True, text_direction=True):
            row.sort_in_line_order()
            instances.extend(row)
        self.reset(instances)


class ElementCollection(Collection):
    """``Element`` 实例的集合。"""

    def _update_bbox(self, e: Element):
        """更新父级边界框。"""
        if not self._parent is None:  # 注意：这里 `if self._parent` 不起作用
            self._parent.union_bbox(e)

    def append(self, e: Element):
        """附加一个实例，相应地更新父级的边界框并设置添加的实例的父级。

        参数：
            e (Element): 要附加的实例。
        """
        if not e:
            return
        self._instances.append(e)
        self._update_bbox(e)

        # 设置父级
        if not self._parent is None:
            e.parent = self._parent

    def insert(self, nth: int, e: Element):
        """插入一个元素并相应地更新父级的边界框。

        参数：
            nth (int): 插入的位置。
            e (Element): 要插入的实例。
        """
        if not e:
            return
        self._instances.insert(nth, e)
        self._update_bbox(e)
        e.parent = self._parent  # 设置父级

    def pop(self, nth: int):
        """删除第 ``nth`` 个实例。

        参数：
            nth (int): 要移除的位置。

        返回：
            Collection: 移除的实例。
        """
        return self._instances.pop(nth)

    def is_flow_layout(self, line_separate_threshold: float, cell_layout=False):
        """判断包含的元素是否为流式布局。"""
        # 如果是垂直文本但不是单元格布局则为浮动布局，因为垂直文本
        # 将使用流表格进行模拟
        if not cell_layout and self.is_vertical_text:
            return False

        # 如果仅有单列则为流式布局
        if len(self) <= 1:
            return True
        if len(self.group_by_columns()) > 1:
            return False

        # 在物理行中分组并检查行间距离
        idx0, idx1 = (0, 2) if self.is_horizontal_text else (3, 1)
        for row in self.group_by_physical_rows(text_direction=True):
            for i in range(1, len(row)):
                dis = abs(row[i].bbox[idx0] - row[i - 1].bbox[idx1])
                if dis >= line_separate_threshold:
                    return False

        return True

    def contained_in_bbox(self, bbox):
        """过滤包含在目标边界框中的实例。

        参数：
            bbox  (fitz.Rect): 目标边界框。
        """
        instances = list(filter(lambda e: bbox.contains(e.bbox), self._instances))
        return self.__class__(instances)

    def split_with_intersection(self, bbox: fitz.Rect, threshold: float = 1e-3):
        """将实例分为两组：一组与 ``bbox`` 相交，另一组不相交。

        参数：
            bbox (fitz.Rect): 目标矩形框。
            threshold (float): 当重叠率超过此阈值时认为相交。默认为0。

        返回：
            tuple: 以原始类类型表示的两组。
        """
        intersections, no_intersections = [], []
        for instance in self._instances:
            # A 包含 B => A & B = B
            intersection = instance.bbox & bbox
            if intersection.is_empty:
                no_intersections.append(instance)
            else:
                factor = round(intersection.get_area() / instance.bbox.get_area(), 2)
                if factor >= threshold:
                    intersections.append(instance)
                else:
                    no_intersections.append(instance)
        return self.__class__(intersections), self.__class__(no_intersections)
